﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace PartialExpressionCompleter
{
    // Originally from http://www.emilstefanov.net/Projects/DisjointSets.aspx
    // MIT licensed by Emil Stefanov.
    class DisjointSets<T>
    {
        /// <summary>
        /// Create an empty DisjointSets data structure
        /// </summary>
        public DisjointSets()
            : this(Enumerable.Empty<T>())
        {
        }

        /// <summary>
        /// Create a DisjointSets data structure with a specified number of elements (with element id's from 0 to count-1)
        /// </summary>
        /// <param name="count"></param>
        public DisjointSets(IEnumerable<T> els)
        {
            m_elementCount = 0;
            m_setCount = 0;
            m_nodes = new List<Node>();
            AddElements(els);
        }

        /// <summary>
        /// Find the set identifier that an element currently belongs to.
        /// Note: some internal data is modified for optimization even though this method is consant.
        /// </summary>
        /// <param name="element"></param>
        /// <returns></returns>
        public int FindElementId(T el)
        {
            return nodeLookup[el].Index;
        }

        public int FindSet(T el)
        {
            return FindSet(FindElementId(el));
        }

        public T FindRoot(T el)
        {
            return m_nodes[FindSet(FindElementId(el))].Data;
        }

        /// <summary>
        /// Find the set identifier that an element currently belongs to.
        /// Note: some internal data is modified for optimization even though this method is consant.
        /// </summary>
        /// <param name="element"></param>
        /// <returns></returns>
        public int FindSet(int elementId)
        {
            if (elementId >= m_elementCount)
                throw new ArgumentOutOfRangeException("elementId");

            Node curNode;

            // Find the root element that represents the set which `elementId` belongs to
            curNode = m_nodes[elementId];
            while (curNode.Parent != null)
                curNode = curNode.Parent;
            Node root = curNode;

            // Walk to the root, updating the parents of `elementId`. Make those elements the direct
            // children of `root`. This optimizes the tree for future FindSet invokations.
            curNode = m_nodes[elementId];
            while (curNode != root)
            {
                Node next = curNode.Parent;
                curNode.Parent = root;
                curNode = next;
            }

            return root.Index;
        }

        public void Union(T el1, T el2)
        {
            Union(FindSet(el1), FindSet(el2));
        }

        /// <summary>
        /// Combine two sets into one. All elements in those two sets will share the same set id that can be gotten using FindSet.
        /// </summary>
        /// <param name="setId1"></param>
        /// <param name="setId2"></param>
        public void Union(int setId1, int setId2)
        {
            if (setId1 >= m_elementCount)
                throw new ArgumentOutOfRangeException("setId1");
            if (setId2 >= m_elementCount)
                throw new ArgumentOutOfRangeException("setId2");

            if (setId1 == setId2)
                return; // already unioned

            Node set1 = m_nodes[setId1];
            Node set2 = m_nodes[setId2];

            // Determine which node representing a set has a higher rank. The node with the higher rank is
            // likely to have a bigger subtree so in order to better balance the tree representing the
            // union, the node with the higher rank is made the parent of the one with the lower rank and
            // not the other way around.
            if (set1.Rank > set2.Rank)
                set2.Parent = set1;
            else if (set1.Rank < set2.Rank)
                set1.Parent = set2;
            else // set1.Rank == set2.Rank
            {
                set2.Parent = set1;
                ++set1.Rank; // update rank
            }

            // Since two sets have fused into one, there is now one less set so update the set count.
            --m_setCount;
        }

        /// <summary>
        /// Add a specified number of elements to the DisjointSets data structure. The element id's of the new elements are numbered
        /// consequitively starting with the first never-before-used elementId.
        /// </summary>
        /// <param name="addCount"></param>
        public void AddElements(IEnumerable<T> els)
        {
            // insert and initialize the specified number of element nodes to the end of the `m_nodes` array
            foreach(var el in els)
            {
                Node newNode = new Node();
                newNode.Parent = null;
                newNode.Index = m_elementCount++;
                newNode.Rank = 0;
                newNode.Data = el;
                m_nodes.Add(newNode);
                nodeLookup.Add(el, newNode);
            }

            // update element and set counts
            m_setCount += els.Count();
        }

        /// <summary>
        /// Returns the number of elements currently in the DisjointSets data structure.
        /// </summary>
        public int ElementCount
        {
            get { return m_elementCount; }
        }

        /// <summary>
        /// Returns the number of sets currently in the DisjointSets data structure.
        /// </summary>
        public int SetCount
        {
            get { return m_setCount; }
        }



        /// <summary>
        /// Internal Node data structure used for representing an element.
        /// </summary>
        private class Node
        {
            /// <summary>
            /// This roughly represent the max height of the node in its subtree.
            /// </summary>
            public int Rank;

            /// <summary>
            /// The index of the element the node represents.
            /// </summary>
            public int Index;

            /// <summary>
            /// The parent node of the node.
            /// </summary>
            public Node Parent;

            public T Data;
        }

        /// <summary>
        /// The number of elements currently in the DisjointSets data structure.
        /// </summary>
        private int m_elementCount;

        /// <summary>
        /// The number of sets currently in the DisjointSets data structure.
        /// </summary>
        private int m_setCount;

        /// <summary>
        /// The list of nodes representing the elements.
        /// </summary>
        private List<Node> m_nodes;

        private Dictionary<T, Node> nodeLookup = new Dictionary<T, Node>();
    }

}
